import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transformes
from torch.utils.data import DataLoader

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1=nn.Conv2d(1,16,3,1,1)
        self.pool=nn.MaxPool2d(2,2)
        self.conv2=nn.Conv2d(16,32,3,1,1)
        self.fc1=nn.Linear(32*7*7,128)
        self.dropout=nn.Dropout(0.5)
        self.fc2=nn.Linear(128,10)
    def forward(self,x):
        x=self.pool(torch.relu(self.conv1(x)))
        x=self.pool(torch.relu(self.conv2(x)))
        x=x.view(-1,32*7*7)
        x=torch.relu(self.fc1(x))
        x=self.dropout(x)
        x=self.fc2(x)
        return x
transform_train=transformes.Compose([
    transformes.RandomCrop(28),
    transformes.RandomHorizontalFlip(),
    transformes.ToTensor(),
    transformes.Normalize(0.5,0.5)
])

transform_test=transformes.Compose([
    transformes.ToTensor(),
    transformes.Normalize(0,1)
])

trainset=torchvision.dataset.MNIST(root="./data",train=True,download=True,transform=transform_train)
trainloader=DataLoader(trainset,batch_size=64,shuffle=True)

testset=torchvision.dastaset.MNIST(root="./data",train=False,download=True,transform=transform_test)
testloader=DataLoader(testset,batch_size=64,shuffle=True)

